# 解决Pydantic版本警告问题
try:
    from pydantic import BaseModel
except ImportError:
    from pydantic.v1 import BaseModel

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
# 导入必要的模块
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_community.embeddings.dashscope import DashScopeEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
import os

# 首先加载和处理文档
loader = PyMuPDFLoader("./deepseek-v3-1-4.pdf")
pages = loader.load_and_split()

# 文档切分
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,
    chunk_overlap=200,
    length_function=len,
    add_start_index=True,
)
texts = text_splitter.create_documents(
    [page.page_content for page in pages[:1]]
)

embeddings = DashScopeEmbeddings(
    model="text-embedding-v1",
    dashscope_api_key="sk-2c036d80d14a443f8769e072bb09a2e1",
)
db = FAISS.from_documents(texts, embeddings)
retriever = db.as_retriever(search_kwargs={"k": 5})

template = """Answer the question based only on the following context:
{context}

Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)

# 使用通义千问模型，通过ChatOpenAI的兼容模式
llm = ChatOpenAI(
    model="qwen-plus",
    api_key="sk-2c036d80d14a443f8769e072bb09a2e1",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)

# Chain
rag_chain = (
        {"question": RunnablePassthrough(), "context": retriever}
        | prompt
        | llm
        | StrOutputParser()
)

res = rag_chain.invoke("deepseek V3有多少参数")
print(res)